
class CompositeRewardCalculator:
    def __init__(self):
        """复合奖励计算器"""
        self.components = {}
        self.weights = {}

    def register_component(self, component_name, component, weight=None):
        """注册奖励组件"""
        self.components[component_name] = component
        if weight is not None:
            self.weights[component_name] = weight
        elif hasattr(component, 'weight'):
            self.weights[component_name] = component.weight
        else:
            self.weights[component_name] = 1.0

        print(f"注册组件: {component_name}，权重: {self.weights[component_name]}")

    def unregister_component(self, component_name):
        """注销奖励组件"""
        if component_name in self.components:
            del self.components[component_name]
            del self.weights[component_name]
            print(f"注销组件: {component_name}")
        else:
            print(f"组件 {component_name} 不存在，无法注销。")

    def set_weight(self, component_name, weight):
        """设置组件权重"""
        if component_name in self.components:
            self.weights[component_name] = weight
            print(f"设置组件 {component_name} 权重为 {weight}")
        else:
            print(f"组件 {component_name} 不存在，无法设置权重。")

    
    def normalize_weights(self):
        """归一化权重"""
        total_weight = sum(self.weights.values())
        if total_weight > 0:
            for component_name in self.weights:
                self.weights[component_name] /= total_weight
        else:
            print("所有组件的权重总和为0,无法归一化。")

    def calculate_reward(self, current_activity, next_activity, case_events, context=None):
        """计算复合奖励"""
        if not self.components:
            return 0.0

        self.normalize_weights()
        reward = 0.0
        rewards_breakdown = {}
        for component_name, component in self.components.items():
            weight = self.weights[component_name]
            component_reward = component.calculate(current_activity, next_activity, case_events, context)
            reward += weight * component_reward
            rewards_breakdown[component_name] = {
                'raw': component_reward,
                'weight': self.weights[component_name],
                'weighted': component_reward*weight
            }
        
         # 如果需要详细分解可以返回分解信息
        if context is not None and isinstance(context, dict):
            context['reward_breakdown'] = rewards_breakdown
            
        return reward